
"""
小说分析
"""
from typing import Union, Optional
from pathlib import Path
from core.parse_novel import parse_novel, merge_chapters
from models import Novel, Chapter
from core.alone_analyze import (
    SimplyExtractor, StyleAnalyzer
)
from core.context_analyze import (
    PlotAnalyzer, CharactersExtractor, CharactersAnalyzer, WorldviewAnalyzer, ProblemsAnalyzer
)
from core.novel_rating import (
    PlotRating, CharacterRating, WorldviewRating, StyleRating, ReaderRating
)
import concurrent.futures
from tqdm import tqdm
from core.utils import result_dir, search_characters
from retry import retry
from loguru import logger
import pickle


class NovelAnalyzer(object):

    def __init__(
            self,
            novel_file: Union[str, Path],
            topn: Optional[int] = 100,
            max_len: Optional[int] = 32000,
            allow_load_existed_result = True,
    ):
        self.analyze_completed = False
        self._load_novel(novel_file, topn, max_len, allow_load_existed_result)
        logger.debug("初始化完成.")

    def _load_novel(
            self,
            novel_file: Union[str, Path],
            topn: Optional[int] = 100,
            max_len: Optional[int] = 32000,
            allow_load_existed_result = True,
    ):
        """加载小说"""
        # 如果结果文件存在，就直接从文件中加载。
        if allow_load_existed_result:
            res_file = result_dir / f'{novel_file.stem}({topn}章).pkl'
            if res_file.exists():
                self.novel = pickle.load(open(res_file, 'rb'))
                self.analyze_completed = True
                return

        logger.info(f"加载小说 {novel_file} ...")
        if isinstance(novel_file, str):
            novel_file = Path(novel_file)
        self.novel = Novel(title=novel_file.stem, topn=topn)
        try:
            novel_texts = open(novel_file, 'r', encoding='GB18030').readlines()
        except:
            novel_texts = open(novel_file, 'r', encoding='utf-8').readlines()
        self.novel.chapters = parse_novel(novel_texts)
        if topn is not None:
            self.novel.chapters = self.novel.chapters[:topn]
        if max_len is not None:
            self.novel.chapters = merge_chapters(self.novel.chapters, max_len)

    def analyze(self):
        """ 小说分析 """
        if self.analyze_completed:
            return self.novel

        self._init_analyze_class()

        # 独立分析
        logger.info("独立分析 ...")
        self.alone_analyze()

        # 上下文分析
        logger.info("上下文分析 ...")
        self.context_analyze()

        # 小说评分
        self.rating_analyze()
        self.reader_rating.score(self.novel)

        # 保存分析结果
        try:
            self.save_novel()
        except Exception as e:
            logger.error(e)

        self.analyze_completed = True

        return self.novel

    def _init_analyze_class(self):
        """
        初始化小说的分析类
        """
        logger.info("初始化分析类 ...")

        # 根据当前章节内容进行分析
        self.simply_extractor = SimplyExtractor()                   # a1. 人物简介抽取
        self.style_analyzer = StyleAnalyzer()                       # a2. 文笔写作分析

        # 根据上下文信息和当前章节内容进行分析。
        self.plot_analyzer = PlotAnalyzer()                         # c1. 情节设计分析
        self.characters_extractor = CharactersExtractor()           # c2. 人物信息抽取
        self.characters_analyzer = CharactersAnalyzer()             # c3. 人物塑造分析
        self.worldview_analyzer = WorldviewAnalyzer()               # c4. 世界观构建分析
        self.problems_analyzer = ProblemsAnalyzer()                 # c5. 问题点分析

        # 小说评分
        self.plot_rating = PlotRating()                             # r1. 情节设计评分
        self.character_rating = CharacterRating()                   # r2. 人物塑造评分
        self.worldview_rating = WorldviewRating()                   # r3. 世界观构建评分
        self.style_rating = StyleRating()                           # r4. 文笔写作评分
        self.reader_rating = ReaderRating()                         # 老白视角的评分计算

    def alone_analyze(self, workers: int = 8):
        """不依赖上下文的分析"""
        tasks = []
        with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as executor:
            for idx, chapter in enumerate(self.novel.chapters):
                # a1. 情节设计分析
                tasks.append(executor.submit(self.simply_extractor.invoke, chapter))
                # a2. 文笔写作分析
                tasks.append(executor.submit(self.style_analyzer.invoke, chapter))

            for future in concurrent.futures.as_completed(tasks):
                if future.exception() is not None:
                    logger.error(future.exception())
                    continue
                logger.info(future.result())

    def context_analyze(self):
        """依赖上下文的分析"""
        for idx, chapter in tqdm(enumerate(self.novel.chapters)):
            tasks = []
            last_chapter = self.novel.chapters[idx - 1] if idx > 0 else Chapter()
            archived_characters = search_characters(self.novel.characters, chapter.character_profile)
            with concurrent.futures.ThreadPoolExecutor() as executor:
                # c1. 情节设计分析
                tasks.append(executor.submit(self.plot_analyzer.invoke, last_chapter, chapter))
                # c2. 人物信息抽取
                tasks.append(executor.submit(self.characters_extractor.invoke, self.novel, archived_characters, chapter))
                # c3. 人物塑造分析
                tasks.append(executor.submit(self.characters_analyzer.invoke, archived_characters, chapter))
                # c4. 世界观构建分析
                tasks.append(executor.submit(self.worldview_analyzer.invoke, last_chapter, chapter))
                # c5. 问题点分析
                tasks.append(executor.submit(self.problems_analyzer.invoke, last_chapter, archived_characters, chapter))

                for future in concurrent.futures.as_completed(tasks):
                    if future.exception() is not None:
                        logger.error(future.exception())
                        continue
                    logger.info(future.result())

    def rating_analyze(self):
        """小说评分"""
        tasks = []
        with concurrent.futures.ThreadPoolExecutor() as executor:
            # r1. 情节设计评分
            tasks.append(executor.submit(self.plot_rating.invoke,self.novel))
            # r2. 人物塑造评分
            tasks.append(executor.submit(self.character_rating.invoke,self.novel))
            # r3. 世界观构建评分
            tasks.append(executor.submit(self.worldview_rating.invoke,self.novel))
            # r4. 文笔写作评分
            tasks.append(executor.submit(self.style_rating.invoke,self.novel))

            for future in concurrent.futures.as_completed(tasks):
                if future.exception() is not None:
                    logger.error(future.exception())
                    continue
                logger.info(future.result())

    @retry(tries=3, delay=0, logger=logger)
    def save_novel(self):
        if not result_dir.exists():
            result_dir.mkdir(parents=True, exist_ok=True)

        with open(result_dir / f'{self.novel.title}({self.novel.topn}章).pkl', 'wb') as fw:
            pickle.dump(self.novel, fw)


if __name__ == '__main__':
    from core.utils.path import data_dir

    for file in data_dir.rglob('*.txt'):
        try:
            analyzer = NovelAnalyzer(file, topn=100, max_len=32000)
            analyzer.analyze()
        except Exception as e:
            logger.error(e)
